import zlib

# 转为字节：每8位一组
def binStrToBytes(s:str):
    s = s.ljust((len(s) + 7) // 8 * 8, '0')  # 补齐到8的倍数
    return bytes([int(s[i:i+8], 2) for i in range(0, len(s), 8)])

# 转回二进制字符串（去除补位需注意长度）
def bytesToBinStr(data: bytes, length: int) -> str:
    """将字节转换为二进制字符串，并截取原始长度"""
    bin_str = ''.join(f'{byte:08b}' for byte in data)
    return bin_str[:length]  # 截取原始长度，去除补零

# 转为字节：每8位一组
def bin_str_to_bytes(s:str):
    s = s.ljust((len(s) + 7) // 8 * 8, '0')  # 补齐到8的倍数
    return bytes([int(s[i:i+8], 2) for i in range(0, len(s), 8)])

# 转为二进制字符串：将字节流还原为 '0' 和 '1' 组成的字符串
def bytes_to_bin_str(data: bytes) -> str:
    return ''.join(f'{byte:08b}' for byte in data)


def compressMethod(lsbs:list[str])->bytes:
    lsbsStr = ''.join(lsbs)  # 使用逗号和空格分隔
    data = binStrToBytes(lsbsStr)
    compressed = zlib.compress(data)
    # print(f"zlib 压缩前: {len(data)} 字节")
    # print(f"zlib 压缩后: {len(compressed)} 字节")
    # print(f"压缩比: {len(compressed)/len(data):.2f}")
    return compressed

# 解压：将压缩数据byte解压还原为bit列表（list[str]）
def decompressMethod(compressed: bytes) -> list[str]:
    decompressed = zlib.decompress(compressed)
    bin_str = bytes_to_bin_str(decompressed)
    return list(bin_str)

if __name__ == "__main__":
    # 示例：0 占 80%，1 占 20%
    lsbs = ["0"] * 800 + ["1"] * 200
    compressed = compressMethod(lsbs)
    print(f"压缩前bytes长度: {len(lsbs)} 压缩后bytes长度: {len(compressed)}")

    originalLength = len(lsbs)
    # 解压测试
    decompressedLsbsBytes = decompressMethod(compressed)
    print(f"解压后bytes长度: {len(decompressedLsbsBytes)}")
    
    decompressed_lsbs_str = bytesToBinStr(decompressedLsbsBytes,originalLength)

    # 验证是否一致
    assert lsbs == list(decompressed_lsbs_str), "解压后数据不一致！"
    print("解压成功，数据一致！")        